🌳 EAGLE‑2 动态草稿树 并行解码 · 置信度驱动 · 动态剪枝

# ============================================================
# EAGLE‑2 动态草稿树构建（核心逻辑）
# ============================================================

def eagle2_generate_tree(
    draft_model,
    hidden_states,          # 目标模型上一轮输出的隐藏状态
    token_ids,              # 当前已生成的 token 序列
    max_nodes = 8,       # 总节点数上限（提交给目标模型验证的草稿数量）
    top_k = 4,            # 每层保留的候选分支数
    max_depth = 4          # 最大树深度
):
    # 1. 初始化：根节点（当前已生成的位置）
    tree = [{
        'tokens': token_ids,
        'hidden': hidden_states,
        'prob': 1.0,           # 路径累积概率
        'depth': 0
    }]

    # 2. 循环扩展，直到达到最大节点数或最大深度
    while len(tree) < max_nodes and tree[-1]['depth'] < max_depth:
        current_nodes = [node for node in tree if node['depth'] == tree[-1]['depth']]
        if not current_nodes:
            break

        # 3. 并行前向：一次 forward 得到所有候选 token 的 logits 和置信度
        logits, confidence = draft_model.forward(current_nodes)

        # 4. 对每个节点，选择 Top‑K 个最高置信度的 token 作为下一层分支
        next_nodes = []
        for idx, node in enumerate(current_nodes):
            top_indices = confidence[idx].topk(top_k).indices.tolist()
            for token_idx in top_indices:
                # 计算路径累积概率（用于后续剪枝）
                path_prob = node['prob'] * confidence[idx][token_idx].item()
                next_nodes.append({
                    'tokens': node['tokens'] + [token_idx],
                    'hidden': logits[idx][token_idx],  # 简化：实际为 hidden 更新
                    'prob': path_prob,
                    'depth': node['depth'] + 1
                })

        # 5. 剪枝：如果总节点数超过 max_nodes，只保留累积概率最高的 max_nodes 条路径
        tree += next_nodes
        if len(tree) > max_nodes:
            tree.sort(key=lambda x: x['prob'], reverse=True)
            tree = tree[:max_nodes]

    # 6. 返回树结构（每个节点是一条候选路径）
    return tree


# ============================================================
# 目标模型验证阶段（简略）
# ============================================================
def verify_with_target(target_model, tree):
    # 将所有候选路径拼接成 batch，并行送入目标模型
    batch = [node['tokens'] for node in tree]
    outputs = target_model.forward(batch)
    # 逐位置判断接受/拒绝（树验证 + 接受率计算）
    accepted_tokens = []
    for i, tokens in enumerate(batch):
        if outputs[i].accepted:   # 该路径被接受
            accepted_tokens.append(tokens[-1])  # 取最后一个 token
    return accepted_tokens